#include "rknn_api.h"
#include <stdio.h>
#include <stdlib.h>

static unsigned char *load_model(const char *filename, int *model_size)
{
    FILE *fp;
    unsigned char *data;

    fp = fopen(filename, "rb");
    if (NULL == fp)
    {
        printf("Open file %s failed.\n", filename);
        return NULL;
    }

    fseek(fp, 0, SEEK_END);
    int size = ftell(fp);

    data = (unsigned char *)malloc(size);
    if (data == NULL)
    {
        printf("buffer malloc failure.\n");
        fclose(fp);
        return NULL;
    }

    fseek(fp, 0, SEEK_SET);
    fread(data, 1, size, fp);
    fclose(fp);

    *model_size = size;
    return data;
}

int main()
{
    const char *model_path = "model/wuzi_640x640.rknn";
    int model_data_size = 0;
    unsigned char *model_data = load_model(model_path, &model_data_size);

    if (model_data == NULL)
    {
        printf("Failed to load model file: %s\n", model_path);
        return -1;
    }

    printf("Model loaded successfully, size: %d bytes\n", model_data_size);

    rknn_context ctx;
    int ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
    if (ret < 0)
    {
        printf("rknn_init error ret=%d\n", ret);
        free(model_data);
        return -1;
    }

    // 查询输入输出数量
    rknn_input_output_num io_num;
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0)
    {
        printf("rknn_query IN_OUT_NUM error ret=%d\n", ret);
        rknn_destroy(ctx);
        free(model_data);
        return -1;
    }
    printf("Model has %d inputs, %d outputs\n", io_num.n_input, io_num.n_output);

    // 查询输入属性
    for (uint32_t i = 0; i < io_num.n_input; i++)
    {
        rknn_tensor_attr attr;
        attr.index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &attr, sizeof(rknn_tensor_attr));
        if (ret < 0)
        {
            printf("rknn_query INPUT_ATTR error ret=%d\n", ret);
            continue;
        }

        printf("Input %d:\n", i);
        printf("  name: %s\n", attr.name);
        printf("  n_dims: %d\n", attr.n_dims);
        printf("  dims: [");
        for (uint32_t j = 0; j < attr.n_dims; j++)
        {
            printf("%d", attr.dims[j]);
            if (j < attr.n_dims - 1)
                printf(", ");
        }
        printf("]\n");
        printf("  fmt: %d (0=NCHW, 1=NHWC)\n", attr.fmt);
        printf("  type: %d\n", attr.type);
        printf("  size: %d\n", attr.size);
    }

    // 查询输出属性
    for (uint32_t i = 0; i < io_num.n_output; i++)
    {
        rknn_tensor_attr attr;
        attr.index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &attr, sizeof(rknn_tensor_attr));
        if (ret < 0)
        {
            printf("rknn_query OUTPUT_ATTR error ret=%d\n", ret);
            continue;
        }

        printf("Output %d:\n", i);
        printf("  name: %s\n", attr.name);
        printf("  n_dims: %d\n", attr.n_dims);
        printf("  dims: [");
        for (uint32_t j = 0; j < attr.n_dims; j++)
        {
            printf("%d", attr.dims[j]);
            if (j < attr.n_dims - 1)
                printf(", ");
        }
        printf("]\n");
        printf("  fmt: %d (0=NCHW, 1=NHWC)\n", attr.fmt);
        printf("  type: %d\n", attr.type);
        printf("  size: %d\n", attr.size);
    }

    rknn_destroy(ctx);
    free(model_data);

    return 0;
}